"""Benchmark: np.full() vs np.ones() * scalar for YoloX preprocessing padding.

Uses `memray run` + `memray stats --json` to measure allocations for both
approaches to creating the padded background array in yolox.preprocess().
Each approach runs N iterations inside a single memray trace, and total
allocations are divided by N to get per-call figures.

Usage:
    python benchmarks/bench_preprocess_padding.py
    python benchmarks/bench_preprocess_padding.py --runs 10
    python benchmarks/bench_preprocess_padding.py --report [PATH]
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
import textwrap
from pathlib import Path

import numpy as np

# YoloX uses (1024, 768) input with 3 channels
SHAPE = (1024, 768, 3)
FILL_VALUE = 114
ARRAY_MIB = np.prod(SHAPE) / (1024 * 1024)

# Script templates — {runs} is substituted at runtime
SCRIPT_TEMPLATES = {
    "np.ones() * 114": textwrap.dedent("""\
        import numpy as np
        for _ in range({runs}):
            _arr = np.ones({shape}, dtype=np.uint8) * {fill}
            del _arr
    """),
    "np.full(114)": textwrap.dedent("""\
        import numpy as np
        for _ in range({runs}):
            _arr = np.full({shape}, {fill}, dtype=np.uint8)
            del _arr
    """),
}

# Patterns that identify benchmark allocations vs import noise
_BENCHMARK_PATTERNS = ("ones", "full", "<module>", "array_multiply", "ones_mul")


def _run_memray(script_body: str) -> dict:
    """Run a script under memray and return JSON stats."""
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(script_body)
        script_path = f.name

    bin_path = tempfile.mktemp(suffix=".bin")
    json_path = tempfile.mktemp(suffix=".json")

    try:
        subprocess.run(
            [sys.executable, "-m", "memray", "run",
             "--trace-python-allocators", "--native", "-o", bin_path, script_path],
            capture_output=True, check=True,
        )
        subprocess.run(
            [sys.executable, "-m", "memray", "stats",
             "--json", "-n", "30", "-o", json_path, bin_path],
            capture_output=True, check=True,
        )
        with open(json_path) as f:
            return json.load(f)
    finally:
        for p in (script_path, bin_path, json_path):
            Path(p).unlink(missing_ok=True)


def _filter_benchmark_allocs(allocs: list[dict]) -> list[dict]:
    return [
        a for a in allocs
        if any(pat in a["location"].split(":")[0] for pat in _BENCHMARK_PATTERNS)
    ]


def _bench_mib(stats: dict, runs: int) -> float:
    """Per-call MiB from benchmark-relevant allocations."""
    total = sum(a["size"] for a in _filter_benchmark_allocs(
        stats["top_allocations_by_size"]
    )) / (1024 * 1024)
    return total / runs


def collect_results(runs: int = 5) -> dict[str, float]:
    """Run each approach in a single memray trace with `runs` iterations."""
    results: dict[str, float] = {}
    for label, template in SCRIPT_TEMPLATES.items():
        script = template.format(runs=runs, shape=SHAPE, fill=FILL_VALUE)
        stats = _run_memray(script)
        results[label] = _bench_mib(stats, runs)
    return results


def print_results(results: dict[str, float], runs: int):
    print(f"\nnp.full() vs np.ones() * scalar")
    print(
        f"Shape: {SHAPE} ({ARRAY_MIB:.2f} MiB)  |  "
        f"{runs} iteration{'s' if runs > 1 else ''}  |  "
        f"Python {sys.version.split()[0]}  |  NumPy {np.__version__}"
    )
    print()

    for label, mib in results.items():
        print(f"  {label}: {mib:.2f} MiB / call")

    ones_mib = results["np.ones() * 114"]
    full_mib = results["np.full(114)"]
    savings = ones_mib - full_mib
    pct = (savings / ones_mib) * 100 if ones_mib > 0 else 0
    print(f"\n  Savings: {savings:.2f} MiB ({pct:.0f}%)")


def generate_report(results: dict[str, float], runs: int, output: str):
    """Generate a comparison chart using Plotly."""
    import plotly.graph_objects as go

    ones_mib = results["np.ones() * 114"]
    full_mib = results["np.full(114)"]
    savings = ones_mib - full_mib
    pct = (savings / ones_mib) * 100 if ones_mib > 0 else 0

    fig = go.Figure(go.Waterfall(
        orientation="v",
        x=["np.ones() * 114", "* scalar temp", "np.full()"],
        y=[ones_mib, -savings, 0],
        measure=["absolute", "relative", "total"],
        text=[f"{ones_mib:.2f} MiB", f"-{savings:.2f} MiB", f"{full_mib:.2f} MiB"],
        textposition="outside", cliponaxis=False,
        textfont=dict(size=14, color="#374151"),
        connector=dict(line=dict(color="#d1d5db", width=1.5, dash="dot")),
        decreasing=dict(marker=dict(color="#f97316")),
        increasing=dict(marker=dict(color="#22c55e")),
        totals=dict(marker=dict(color="#6366f1")),
    ))

    fig.add_annotation(
        x="np.full()", y=full_mib * 0.5,
        text=f"<b>-{pct:.0f}%</b>",
        showarrow=False, font=dict(size=18, color="white"),
    )

    subtitle = (
        f"YoloX preprocess  |  {SHAPE} = {ARRAY_MIB:.2f} MiB  |  "
        f"{runs} iteration{'s' if runs > 1 else ''}  |  "
        f"Python {sys.version.split()[0]}  |  NumPy {np.__version__}"
    )

    fig.update_layout(
        template="simple_white",
        paper_bgcolor="white", plot_bgcolor="white",
        font=dict(family="Inter, sans-serif", color="#374151"),
        title=dict(
            text=(
                f"<b>np.full() vs np.ones() * scalar</b>"
                f"<br><span style='font-size:11px;color:#9ca3af'>{subtitle}</span>"
            ),
            x=0.5, xanchor="center", font=dict(size=16),
        ),
        showlegend=False,
        yaxis=dict(title="Memory (MiB)", range=[0, ones_mib * 1.35], gridcolor="#f3f4f6"),
        margin=dict(l=60, r=40, t=90, b=50),
        height=420, width=580,
    )

    fig.write_image(output, scale=2)
    print(f"  Report: {output}")


def main():
    parser = argparse.ArgumentParser(description="Benchmark np.full vs np.ones * scalar")
    parser.add_argument(
        "--runs", type=int, default=5,
        help="Number of iterations per approach (default: 5)",
    )
    parser.add_argument(
        "--report", nargs="?", const="benchmarks/bench_preprocess_padding.png",
        metavar="PATH", help="Generate chart (default: benchmarks/bench_preprocess_padding.png)",
    )
    args = parser.parse_args()

    results = collect_results(runs=args.runs)
    print_results(results, runs=args.runs)
    if args.report is not None:
        generate_report(results, runs=args.runs, output=args.report)


if __name__ == "__main__":
    main()